import os

def delete_init_files(directory, excluded_dirs=[]):
    """递归删除目录及其子目录中的所有__init__.py文件"""
    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        if item not in excluded_dirs:
            if os.path.isdir(item_path) and not item.startswith(('__', '.')):
                delete_init_files(item_path, excluded_dirs = excluded_dirs)  # 递归删除子目录中的__init__.py
            elif item == '__init__.py':
                os.remove(item_path)  # 删除找到的__init__.py文件

def generate_init_file(directory, root_package='', excluded_dirs=[], excluded_dirs2=[]):
    """生成或重新生成__init__.py文件，包含导入语句，但排除在excluded_dirs和excluded_dirs2中的目录"""
    # 如果当前目录在excluded_dirs2中，则不执行任何操作
    if directory.split(os.sep)[-1] in excluded_dirs2:
        return

    # 首先递归删除所有现有的__init__.py文件
    delete_init_files(directory)
    
    # __init__.py文件的路径
    init_file_path = os.path.join(directory, '__init__.py')

    # 创建并写入新的__init__.py文件
    with open(init_file_path, 'w') as init_file:
        # 遍历目录下的所有文件和子目录
        for item in os.listdir(directory):
            item_path = os.path.join(directory, item)
            
            # 如果是Python文件，并且不是__init__.py，同时目录不在excluded_dirs中
            if item.endswith('.py') and item != '__init__.py' and directory.split(os.sep)[-1] not in excluded_dirs:
                module_name = os.path.splitext(item)[0]
                init_file.write(f'from .{module_name} import *\n')
            
            # 如果是子目录item，并且不是以'__'或'.'开头，同时不在excluded_dirs2中
            elif os.path.isdir(item_path) and not item.startswith(('__', '.')) and item not in excluded_dirs2:
                # 递归生成子目录的__init__.py
                sub_package_import = f'{root_package}.{item}' if root_package else item
                if directory.split(os.sep)[-1] in excluded_dirs:
                    excluded_dirs.append(item)
                generate_init_file(item_path, root_package=sub_package_import, excluded_dirs=excluded_dirs, excluded_dirs2=excluded_dirs2)
                if directory.split(os.sep)[-1] not in excluded_dirs:
                    init_file.write(f'from .{item} import *\n')
            

# 调用generate_init_file函数，生成或重新生成__init__.py文件
generate_init_file('Training01_whl', excluded_dirs=['utils'], excluded_dirs2=['data'])